/*
 * This tool make the original data set
 * The data set path is given as a command
 * line parameter to this bin.
 * Than Images will be transformed to lmdb
 * database.
 * The path of database is also given as
 * a commandline parameter
 *
 * The database created by this file is
 * used as the overall database. which
 * contains all interested pixels.
 */

#include "tools.h"

using google::INFO;
using google::WARNING;
using google::FATAL;
using google::ERROR;
using boost::filesystem::path;
using boost::filesystem::exists;

/*
 * The file list should be a table in csv format;
 * any blank space will be trimmed.
 *
 */
DEFINE_string(filelist,
              "/home/d/workspace/DeepDR/data/imgs_hospital/trainfiles.csv",
              "The image list in csv format");

DEFINE_string(outDB,
              "overAllDataSet",
              "The directory that lmdb file is writen to");

DEFINE_string(minorityDB,
              "minority",
              "the minority samples classes");

DEFINE_bool(overriteDB,
            false,
            "Whether overwrite original database if exist"
);

int main(int argc, char **argv) {
    google::SetCommandLineOption("alsologtostderr", "true");
    google::SetCommandLineOption("colorlogtostderr", "true");
    google::InitGoogleLogging(argv[0]);
    google::SetUsageMessage(
            "Convert the original image set to lmdb database\n"
    );
    google::ParseCommandLineFlags(&argc, &argv, true);
    LOG(INFO) << "Info";
    LOG(WARNING) << "Warning";
    LOG(INFO) << "Reading Image set :" << FLAGS_filelist;
    LOG(INFO) << "Writing to :" << FLAGS_outDB;
    backward::SignalHandling sh;
    if (!sh.loaded()) {
        LOG(ERROR) << "BACKWARD NOT LOADED";
    }

    LOG(INFO) << "Building original database";
    if (exists(path(FLAGS_outDB)) && !FLAGS_overriteDB) {
        LOG(ERROR) << "Error, database " <<
        FLAGS_outDB <<
        " already exists";
    } else {
        ImgDataGenerator dataGenerator(FLAGS_filelist);
        LOG(INFO) << "Writing " << dataGenerator.length() << "Objects";
        if (exists(path(FLAGS_outDB))) {
            LOG(WARNING) << "overwriting original database";
            boost::filesystem::remove_all(path(FLAGS_outDB));
        }
        LMDBDataSet lmdbDataSet;
        lmdbDataSet.newDB(FLAGS_outDB);
        size_t displayScale = dataGenerator.length() / 100;
        for (size_t i = 0; i < dataGenerator.length(); i++) {
            lmdbDataSet.addSample(dataGenerator[i]);
            if (i % displayScale == 0) {
                LOG(INFO) << "finished " <<
                double(i) / double(dataGenerator.length()) * 100 <<
                "%";
            }
        }
        lmdbDataSet.close();
    }

    LOG(INFO) << "Building database for minority classes";
    if (exists(path(FLAGS_minorityDB)) && !FLAGS_overriteDB) {
        LOG(ERROR) << "Error, database " <<
        FLAGS_minorityDB <<
        " already exists";
    } else {
        if (exists(path(FLAGS_minorityDB))) {
            LOG(WARNING) << "overwriting minority database";
            boost::filesystem::remove_all(path(FLAGS_outDB));
        }
        caffe::Datum datum;
        LMDBDataSet minority, original;
        minority.newDB(FLAGS_minorityDB);
        original.loadFromLMDB(FLAGS_outDB);
        int key = 0;
        for (int i = 0; i < original.length(); i++) {
            original.seek(i);
            original.getDatum(datum);
            if (datum.label() > 0) {
                minority.addDatum(key++, &datum);
            }
            if (i % 10000 == 0) {
                LOG(INFO) << double(i) / original.length() * 100 << "%";
            }
        }
        minority.close();
        original.close();
    }
}
